from openai import OpenAI
from typing import Dict, Any, List
import json
import os

# 预定义文章类别
ARTICLE_CATEGORIES = [
    "科技", "教育", "健康", "财经", "体育",
    "文化", "娱乐", "政治", "环境", "社会"
]

# 文章分类专家提示词
ARTICLE_CLASSIFIER_PROMPT = f"""你是文章分类专家。职责：
分析文章内容，将文章分类到预定义类别中（可以是单个或多个类别），输出格式如下：
{{
    "主要类别": "",  // 最主要的一个类别
    "次要类别": [],  // 其他相关类别（0-2个）
    "分类依据": {{   // 对每个选择的类别说明分类依据
        "主要类别依据": "",
        "次要类别依据": []
    }},
    "关键主题词": [], // 3-5个体现文章主题的关键词或短语
    "置信度": {{      // 对分类结果的置信程度（0-1）
        "主要类别置信度": 0.0,
        "次要类别置信度": []
    }}
}}

预定义类别：
{', '.join(ARTICLE_CATEGORIES)}

注意事项：
1. 主要类别必须且只能选择一个最贴切的类别
2. 次要类别可以选择0-2个相关类别
3. 详细说明每个类别的分类依据
4. 给出分类的置信度评分（0-1之间）
"""

class ArticleClassifier:
    def __init__(self, model_name: str = "qwen-2.5-72b"):
        """初始化配置"""
        self.client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY", "EMPTY"),
            base_url="http://172.16.8.36:9997/v1"
        )
        self.model_name = model_name

    def classify_article(self, article_text: str) -> Dict[str, Any]:
        """对文章进行分类"""
        try:
            # 发送分类请求
            classify_response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": ARTICLE_CLASSIFIER_PROMPT},
                    {"role": "user", "content": article_text}
                ],
                temperature=0.1
            )
            
            classify_data = classify_response.choices[0].message.content
            if not classify_data:
                return None
                
            try:
                json_start = classify_data.find('{')
                json_end = classify_data.rfind('}') + 1
                if json_start >= 0 and json_end > json_start:
                    result = json.loads(classify_data[json_start:json_end])
                else:
                    return None
            except json.JSONDecodeError:
                print("分类结果不是有效的JSON格式")
                return None

            return result

        except Exception as e:
            print(f"处理文章内容时出错: {e}")
            return None

    def save_result(self, result: Dict[str, Any], output_dir: str, file_id: str) -> bool:
        """保存处理结果"""
        try:
            os.makedirs(output_dir, exist_ok=True)
            output_file = os.path.join(output_dir, f"{file_id}.json")
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存结果失败: {e}")
            return False

# 测试代码
if __name__ == "__main__":
    # 测试文章
    test_article = """
    近日，中国空间站迎来了首次商业航天任务。商业航天公司成功将多个科研项目送入太空，
    这些项目包括新型材料实验和生物医学研究。这标志着中国航天产业开始向商业化转型，
    为民营企业参与航天活动开辟了新途径。专家表示，商业航天不仅能推动航天技术创新，
    还将带动相关产业发展，创造可观的经济效益。
    """

    # 初始化分类器
    classifier = ArticleClassifier()
    
    # 处理文章
    result = classifier.classify_article(test_article)
    if result:
        print("\n分类结果:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
        
        # 保存结果
        classifier.save_result(result, "classified_articles", "test_article") 